#!/usr/bin/env python3
"""
Qwen3-4B-AWQ模型测试运行脚本

此脚本提供了一个简单的命令行界面，用于运行Qwen模型的各项测试。
"""

import sys
import os
import argparse
import time
from typing import Dict, Any

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from qwen_model_client import QwenModelClient, ModelConfig
from test_qwen_model import run_performance_benchmark


def interactive_test():
    """交互式测试模式"""
    print("="*60)
    print("Qwen3-4B-AWQ模型交互式测试")
    print("="*60)
    
    # 初始化客户端
    config = ModelConfig(
        host="192.168.1.236",
        port=8000,
        timeout=30
    )
    client = QwenModelClient(config)
    
    # 健康检查
    print("正在进行健康检查...")
    if not client.health_check():
        print("❌ 健康检查失败，请确保模型服务正在运行")
        return
    
    print("✅ 健康检查通过")
    
    # 获取可用模型列表
    try:
        models = client.get_available_models()
        print(f"可用模型: {models}")
    except Exception as e:
        print(f"获取模型列表失败: {str(e)}")
    
    print("\n您可以输入提示与模型交互，输入'quit'退出")
    print("-" * 60)
    
    while True:
        try:
            user_input = input("\n您: ")
            if user_input.lower() in ['quit', 'exit', 'q']:
                print("退出交互模式")
                break
            
            print("模型: ", end="", flush=True)
            start_time = time.time()
            response = client.simple_chat(user_input)
            end_time = time.time()
            
            print(response)
            print(f"[响应时间: {end_time - start_time:.2f}秒]")
            
        except KeyboardInterrupt:
            print("\n\n检测到中断，退出交互模式")
            break
        except Exception as e:
            print(f"\n错误: {str(e)}")


def quick_test():
    """快速测试模式"""
    print("="*60)
    print("Qwen3-4B-AWQ模型快速测试")
    print("="*60)
    
    # 初始化客户端
    config = ModelConfig(
        host="192.168.1.236",
        port=8000,
        timeout=30
    )
    client = QwenModelClient(config)
    
    # 健康检查
    print("正在进行健康检查...")
    if not client.health_check():
        print("❌ 健康检查失败，请确保模型服务正在运行")
        return
    
    print("✅ 健康检查通过")
    
    # 获取可用模型列表
    try:
        models = client.get_available_models()
        print(f"可用模型: {models}")
    except Exception as e:
        print(f"获取模型列表失败: {str(e)}")
    
    # 测试用例
    test_cases = [
        "你好，请介绍一下你自己。",
        "请解释什么是人工智能。",
        "写一个Python函数计算斐波那契数列的第n项。"
    ]
    
    print(f"正在进行 {len(test_cases)} 项快速测试...")
    
    success_count = 0
    total_time = 0
    
    for i, prompt in enumerate(test_cases):
        print(f"\n测试 {i+1}/{len(test_cases)}: {prompt}")
        try:
            start_time = time.time()
            response = client.simple_chat(prompt)
            end_time = time.time()
            
            response_time = end_time - start_time
            total_time += response_time
            
            print(f"✅ 成功 (耗时: {response_time:.2f}秒)")
            print(f"回复: {response[:100]}..." if len(response) > 100 else f"回复: {response}")
            
            success_count += 1
            
        except Exception as e:
            print(f"❌ 失败: {str(e)}")
    
    print("\n" + "="*60)
    print("测试结果汇总")
    print("="*60)
    print(f"成功: {success_count}/{len(test_cases)}")
    print(f"失败: {len(test_cases) - success_count}/{len(test_cases)}")
    print(f"总耗时: {total_time:.2f}秒")
    if success_count > 0:
        print(f"平均响应时间: {total_time/success_count:.2f}秒")
    
    if success_count == len(test_cases):
        print("\n🎉 所有测试通过！")
    else:
        print(f"\n⚠️ 有 {len(test_cases) - success_count} 项测试失败")


def run_unit_tests():
    """运行单元测试"""
    print("="*60)
    print("运行单元测试")
    print("="*60)
    
    try:
        import unittest
        from test_qwen_model import TestQwenModelClient, TestErrorHandling
        
        # 创建测试套件
        suite = unittest.TestSuite()
        
        # 添加测试用例
        suite.addTest(unittest.makeSuite(TestQwenModelClient))
        suite.addTest(unittest.makeSuite(TestErrorHandling))
        
        # 运行测试
        runner = unittest.TextTestRunner(verbosity=2)
        result = runner.run(suite)
        
        # 输出测试结果
        print("\n" + "="*60)
        print("单元测试结果")
        print("="*60)
        print(f"运行测试: {result.testsRun}")
        print(f"失败: {len(result.failures)}")
        print(f"错误: {len(result.errors)}")
        
        if result.failures:
            print("\n失败的测试:")
            for test, traceback in result.failures:
                print(f"- {test}: {traceback}")
        
        if result.errors:
            print("\n错误的测试:")
            for test, traceback in result.errors:
                print(f"- {test}: {traceback}")
        
        if result.wasSuccessful():
            print("\n✅ 所有单元测试通过！")
        else:
            print("\n❌ 部分单元测试失败")
            
    except Exception as e:
        print(f"运行单元测试时发生错误: {str(e)}")


def run_speed_comparison(host='192.168.1.236', port=8000):
    """运行速度对比测试"""
    print("🚀 运行速度对比测试...")
    
    try:
        # 导入速度对比测试模块
        import speed_comparison
        
        # 运行速度对比测试
        speed_comparison.run_comparison()
        
    except ImportError:
        print("❌ 无法导入speed_comparison模块")
    except Exception as e:
        print(f"❌ 速度对比测试失败: {e}")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="Qwen3-4B-AWQ模型测试工具")
    parser.add_argument(
        "mode", 
        choices=["interactive", "quick", "unit", "benchmark", "speed"],
        nargs="?",
        default="interactive",
        help="测试模式: interactive(交互式), quick(快速测试), unit(单元测试), benchmark(性能基准测试), speed(速度对比测试)"
    )
    
    args = parser.parse_args()
    
    if args.mode == "interactive":
        interactive_test()
    elif args.mode == "quick":
        quick_test()
    elif args.mode == "unit":
        run_unit_tests()
    elif args.mode == "benchmark":
        run_performance_benchmark()
    elif args.mode == "speed":
        run_speed_comparison()


if __name__ == "__main__":
    main()